import glob
import random
import torch
from torch.utils.data import dataset, dataloader
from PIL import Image
from torchvision import transforms
import torch.nn as nn
import torch.nn.functional as F

from nlp.config import Config

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class NlpModule(nn.Module):
    def __init__(self, config):
        super(NlpModule, self).__init__()
        self.embedding = nn.Embedding(config.n_vocab, config.embed_size, padding_idx=config.n_vocab - 1)
        self.lstm = nn.LSTM(config.embed_size, config.hidden_size, config.num_layers, bidirectional=True,
                            batch_first=True, dropout=config.dropout)
        self.maxpool = nn.MaxPool1d(config.pad_size)
        self.fc = nn.Linear(config.hidden_size * 2 + config.embed_size, config.num_classes)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        embed = self.embedding(x)  # [batch_size,seqlen,embed_size]
        out, _ = self.lstm(embed)
        out = torch.cat((embed, out), 2)
        out = F.relu(out)
        out = out.permute(0, 2, 1)
        out = self.maxpool(out)
        out = out.reshape(out.size()[0], -1)
        out = self.fc(out)
        out = self.softmax(out)
        return out


if __name__ == '__main__':
    conf = Config()
    conf.pad_size = 640
    m = NlpModule(conf)

    input_t = torch.tensor([i for i in range(640)]).reshape([1, 640])
    out = m.forward(input_t)
    print(out.size())
